import time
import sys
import os
import azure.cognitiveservices.speech as speechsdk
from concurrent.futures import ProcessPoolExecutor
from functools import partial
import csv
import time
import argparse


def call_ms_service(text, fpath, voice_name):
    api_key = os.environ.get('AZURE_TTS_KEY')
    ms_api_region = os.environ.get('AZURE_TTS_KEY_REGION')
    # print(api_key)
    speech_config = speechsdk.SpeechConfig(subscription=api_key, region=ms_api_region)
    speech_config.set_speech_synthesis_output_format(speechsdk.SpeechSynthesisOutputFormat.Riff24Khz16BitMonoPcm)
    speech_config.speech_synthesis_language = "en-US"
    speech_config.speech_synthesis_voice_name = voice_name
    audio_config = speechsdk.audio.AudioOutputConfig(use_default_speaker=True)
    audio_config = speechsdk.audio.AudioOutputConfig(filename=fpath)
    synthesizer = speechsdk.SpeechSynthesizer(speech_config=speech_config, audio_config=audio_config)
    try:
        result = synthesizer.speak_text(text)
        # Check the result and print details in case of cancellation
        if result.reason == speechsdk.ResultReason.Canceled:
            cancellation_details = result.cancellation_details
            print(f"Speech synthesis canceled: {cancellation_details.reason}")

            if cancellation_details.reason == speechsdk.CancellationReason.Error:
                print(f"Error Code: {cancellation_details.error_code}")
                print(f"Error Details: {cancellation_details.error_details}")
            return None

        print('syn:', text, fpath)
        # time.sleep(1)
        return text
    except Exception as e:
        print(f"An error occurred: {e}")
        return None
    
    
def process(in_txt, out_base_dir):
    # in_txt = './tpo_kmf.csv'
    # out_base_dir = './generated_wav'
    n_jobs = 4
    voice_names = [
                'en-US-MonicaNeural',
                'en-US-RogerNeural',
                'en-US-JennyNeural',
                'en-US-GuyNeural'
                ]
    os.makedirs(out_base_dir, exist_ok=True)
    executor = ProcessPoolExecutor(max_workers=n_jobs)
    futures = []
    for voice_name in voice_names:
        spk = voice_name[len('en-US-'):]
        spk = spk[:-len('Neural')]
        print(spk)
        out_spk_dir = os.path.join(out_base_dir, spk)
        os.makedirs(out_spk_dir, exist_ok=True)
        with open(in_txt, 'r', newline='', encoding='utf-8') as csvfile:
            reader = csv.reader(csvfile, delimiter=',', quotechar='"')
            
            for row in reader:
                if len(row) == 5:
                    p_id = row[0]
                    text = row[1]
                    print(p_id, text)
                    p_id = p_id.strip()
                    text = text.strip()
                    print(spk, p_id, text)
                    f_path = os.path.join(out_spk_dir, '{}.wav'.format(p_id))
                    futures.append(executor.submit(partial(call_ms_service, text, f_path, voice_name)))
        
        print("finished " + spk)
    results = [future.result() for future in futures if future.result() is not None]
    print("finished")
                
def parse_arguments():
    parser = argparse.ArgumentParser(description="Generate speech from text using Azure Cognitive Services.")
    parser.add_argument('-i', '--in_txt', type=str, required=False, default='./tpo_kmf.csv', help='Path to the input CSV file. Default is "./tpo_kmf.csv".')
    parser.add_argument('-o', '--out_dir', type=str, required=False, default='./generated_wav', help='Path to the output directory where WAV files will be saved. Default is "./generated_wav".')
    return parser.parse_args()

if __name__ == '__main__':
    args = parse_arguments()
    process(args.in_txt, args.out_dir)
